import os
import argparse
import re
from sqllineage.runner import LineageRunner

def extract_dml_sql_from_dir(directory):
    """
    遍历指定目录下的所有.pkb和.pks文件，提取其中的DML SQL、CREATE TABLE语句及其所属的存储过程名称。
    :param directory: 要搜索的目录路径。
    """
    all_statements = []
    # 遍历目录下所有.pkb和.pks文件
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.endswith(('.pkb', '.pks')):
                file_path = os.path.join(root, file)
                with open(file_path, 'r', encoding='utf-8') as f:
                    content = f.read()
                    statements_with_proc = extract_dml_sql(content)
                    all_statements.extend(statements_with_proc)
    return all_statements


def extract_dml_sql(content):
    """
    从给定的内容中提取DML SQL、CREATE TABLE语句及其所属的存储过程名称。
    :param content: 包含SQL或PL/SQL代码的字符串。
    :return: 提取出的SQL语句列表，每个元素是一个包含'statement'和'procedure_name'键的字典。
    """
    procedure_name = None
    statements_with_proc = []

    # 移除多行注释 /* ... */
    content = re.sub(r'/\*.*?\*/', '', content, flags=re.DOTALL)
    # 移除单行注释 --
    content = re.sub(r'--.*$', '', content, flags=re.MULTILINE)

    # 正则匹配存储过程声明
    proc_pattern = re.compile(r'PROCEDURE\s+(\w+)\s*\(', re.IGNORECASE)
    # 正则匹配DML语句（INSERT/UPDATE/DELETE/MERGE）、动态SQL以及CREATE TABLE语句
    dml_pattern = re.compile(r'''
        (                             # 捕获组1：静态DML或CREATE TABLE
            (CREATE\s+TABLE|INSERT\s+INTO|UPDATE|DELETE\s+FROM|MERGE\s+INTO)  # 关键字
            .*?                       # 任意字符（非贪婪）
            (?=\b(?:END|EXCEPTION|BEGIN|ELSIF|ELSE|WHEN)\b|;)  # 前瞻结束符
        )
        |
        (                             # 捕获组2：动态SQL
            EXECUTE\s+IMMEDIATE\s+     # EXECUTE IMMEDIATE
            ('(?:''|[^'])*')           # 捕获单引号字符串（处理转义）
        )
    ''', re.IGNORECASE | re.DOTALL | re.VERBOSE)

    # 查找所有的存储过程名称
    for match in proc_pattern.finditer(content):
        procedure_name = match.group(1)

    # 查找所有的DML和CREATE TABLE语句
    for match in dml_pattern.finditer(content):
        static_dml_or_create_table = match.group(1)
        dynamic_sql = match.group(3)

        if static_dml_or_create_table:
            statement = static_dml_or_create_table.strip()
            statement = re.sub(r'\s+', ' ', statement)  # 简化空格
            statements_with_proc.append({'statement': statement + ';', 'procedure_name': procedure_name})
        elif dynamic_sql:
            sql = dynamic_sql.strip("'")
            sql = re.sub(r"''", "'", sql)  # 还原转义的单引号
            if re.match(r'^\s*(INSERT|UPDATE|DELETE|MERGE|CREATE)', sql, re.I):
                statements_with_proc.append({'statement': sql + ';', 'procedure_name': procedure_name})

    return statements_with_proc

def is_valid_statement(sql):
    """
    判断给定的SQL语句是否为有效的。
    如果是UPDATE_ETL开头的函数，则match_result返回True，其他正常的SQL返回False。
    :param sql: 要检查的SQL语句。
    :return: 如果语句有效则返回True，否则返回False。
    """
    # 忽略大小写并检查是否以'UPDATE_ETL'开头
    match_result = sql.strip().upper().startswith('UPDATE_ETL')
    print(f'bRet={match_result}')
    return match_result

def analyze_lineage(statements):
    """
    使用sqllineage分析给定SQL语句的数据血缘关系。
    :param statements: 包含存储过程名称和对应SQL语句的列表。
    """
    for statement_info in statements:
        proc_name = statement_info['procedure_name']
        sql = statement_info['statement']

        # 过滤掉无效的SQL语句
        if  is_valid_statement(sql):
            print(f'---sql={sql}---');
            continue;

        print(f"Analyzing lineage for procedure: {proc_name}->{sql}")
        runner = LineageRunner(sql)

        # 输出源表
        print("Source Tables:")
        for src in runner.source_tables:
            #print(f'---{src}----');
            print(src)

        # 输出目标表
        print("Target Tables:")
        for tgt in runner.target_tables:
            #print(f'---{tgt}---');
            print(tgt)

        print("Intermediate Tables:")
        for intermediate in runner.intermediate_tables:
            #print(f'---{intermediate}---');
            print(intermediate)

        print("\n")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="oracle package analysis");
    parser.add_argument("--out_dir", type=str, default="out");
    parser.add_argument("--data_path", type=str, default="E:/pyworkspace/sqllineage-master/test_data_path");
    args = parser.parse_args();

    # 文件路径
    output_file_path = 'E:/pyworkspace/sqllineage-master/out/allsql.txt';
    # 使用示例
    directory_path = args.data_path;
    statements = extract_dml_sql_from_dir(directory_path)
    with open(output_file_path, 'w', encoding='utf-8') as f:
        for statement_info in statements:
            # 构建要写入的内容
            line = f"Procedure: {statement_info['procedure_name']}, Statement: {statement_info['statement']}\n"
            print(f"Procedure: {statement_info['procedure_name']}, Statement: {statement_info['statement']}");
            # 写入内容到文件
            f.write(line)
            #下面还需要调用血缘分析的函数展示血缘关系图。。。。。。
            #analyze_lineage(statements)